from pydantic import BaseModel
from typing import Dict, List, Union

from src.schemas.category import CategoryReturn
from src.schemas.entry import EntryReturn
from src.schemas.history import HistoryReturn

ListReturn = List[
    Union[
        EntryReturn,
        HistoryReturn,
        CategoryReturn
    ]
]


class ResponseModel(BaseModel):
    code: int
    message: str
    data: Union[
        ListReturn,
        CategoryReturn,
        EntryReturn,
        Dict[str, ListReturn],
        None
    ]

    @staticmethod
    def ReturnSuccess(ret: Union[dict, list, None] = None):
        return ResponseModel(code=0, message="success", data=ret)

    @staticmethod
    def ReturnFailed(message: str, ret: Union[dict, list, None] = None):
        return ResponseModel(code=400, message=message, ret=ret)

    @staticmethod
    def ReturnException(e: Exception):
        return ResponseModel(code=500, message="Exception -> {}".format(str(e)))
